#include "fitness_functions.h"
#include "params.h"
#include "pso.h"
#include <cmath>
#include <iostream>


namespace fitness
{
    // The available fitness functions.
    double fitness_sphere(const Position *p_position);
    double fitness_griewank(const Position *p_position);
    double fitness_rastrigin(const Position *p_position);
    double fitness_pso(const Position *p_position);

    fitness_function getFitnessFunction(unsigned int p_function) 
    {
        switch(p_function) {
          case 1:
            return fitness_sphere;
            break;
          case 2:
            return fitness_griewank;
            break;
          case 3:
            return fitness_rastrigin;
            break;
          case 4:
            return fitness_pso;
            break;
        }

        return NULL;
    }

    double fitness_sphere(const Position *p_position)
    {
        double fitness = 0;
        unsigned int d = p_position->getComponentsSize();

        for(unsigned int i = 0; i < d; ++i) {
            fitness += std::pow(p_position->getComponent(i), 2);
            //std::cout << p_position->getComponent(i) << " ";
        }

        // std::cout << fitness << std::endl;

        return fitness;
    }

    double fitness_griewank(const Position *p_position)
    {
        double fitness;
        double sum = 0;
        double product = 1;
        unsigned int d = p_position->getComponentsSize();

        for(unsigned int i = 0; i < d; ++i) {
            sum += std::pow(p_position->getComponent(i),2);
            product *= std::cos(p_position->getComponent(i) / std::sqrt(i+1));
        }

        fitness = (sum / 4000.0) - product + 1;

        return fitness;
    }

    double fitness_rastrigin(const Position *p_position) 
    {
        double fitness = 0.0;
        unsigned int d = p_position->getComponentsSize();
        double x;

        for(unsigned int i = 0; i < d; ++i) {
            x = p_position->getComponent(i);
            fitness += std::pow(x, 2) - 10*std::cos(2*M_PI*x) + 10;
        }

        return fitness;
    }

    double fitness_pso(const Position *p_position)
    {
        const unsigned int pso_tests = 10;
        Params params;
        double meta_fitness = 0.0;
        double neighbours;
        double neighbour_penalty;
        const unsigned int swarm_size = 30;

        // Calculate how many neighbours to use.
        neighbours = std::modf(p_position->getComponent(4), &neighbour_penalty);
        neighbours *= static_cast<double>(swarm_size);
        
        // Some parameters are now defined by our "position".
        params.alpha_1 = p_position->getComponent(0);
        params.alpha_2 = p_position->getComponent(1);
        params.alpha_3 = p_position->getComponent(2);
        params.omega_bottom = p_position->getComponent(3);
        params.neighbours = std::floor(neighbours + 0.5);
        params.omega_top = p_position->getComponent(5);

        // These are some other default params.
        params.swarm_size = 30;
        params.iterations = 500;
        params.dimensions = 5;
        //params.omega_top = 0.729;
        params.boundary_radius = 6;
        params.fitness_function = 3; // Rastrigin function.

        // Pso with custom params that will not log info.
        Pso pso(&params, false);

        // Obtain the meta-fitness of this pso.
        for(unsigned int i = 0; i < pso_tests; ++i) {
            meta_fitness += pso.exec();
        }
        meta_fitness /= pso_tests;

        // Add penalizations to the meta-fitness for undesirable variables.
        // Penalize values of neighbours above 1.
        meta_fitness += std::pow(std::fabs(neighbour_penalty), 2);
        // Penalize big numbers for alphas.
        meta_fitness += std::pow(10 * std::fabs(params.alpha_1)/35,2) + 
            std::pow(10 * std::fabs(params.alpha_2)/35,2) + 
            std::pow(10 * std::fabs(params.alpha_3)/35,2);

        return meta_fitness;
    }
}
